import json
import tkinter as tk
from tkinter import ttk
import logging
import os
import sys
import uuid
from typing import Dict

from llama_index.core import StorageContext, SimpleDirectoryReader, VectorStoreIndex, Settings, Document
from llama_index.core.callbacks import LlamaDebugHandler, CallbackManager
from llama_index.core.storage.storage_context import DEFAULT_PERSIST_DIR
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.openai.utils import ALL_AVAILABLE_MODELS, CHAT_MODELS
from llama_index.vector_stores.milvus import MilvusVectorStore

from utils import project_setting
from utils.str_utils import str_prominent, format_filename

logger_level = logging.DEBUG
# logger_level = logging.ERROR
logging.basicConfig(stream=sys.stdout, level=logger_level,
                    format='%(asctime)s %(name)s [%(pathname)s line:%(lineno)d] %(levelname)s %(message)s')
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

llama_debug = LlamaDebugHandler(print_trace_on_end=True)
from langfuse.llama_index import LlamaIndexCallbackHandler

langfuse_callback_handler = LlamaIndexCallbackHandler(
    trace_name='demo',
    user_id='sks',
    session_id=str(uuid.uuid1()),
    public_key=os.getenv('LANGFUSE_PUBLIC_KEY'),
    secret_key=os.getenv('LANGFUSE_SECRET_KEY'),
    host="https://us.cloud.langfuse.com"
)

Settings.callback_manager = CallbackManager([langfuse_callback_handler, llama_debug])
# Settings.callback_manager = CallbackManager([langfuse_callback_handler])

# 定义模型及其对应的上下⽂⼤⼩
LLM_MODELS_CONF: Dict[str, int] = {
    "ishumilin/deepseek-r1-coder-tools:14b": 64000,  #设置上下⽂⼤⼩为64000
    "qwen2.5-coder:7b": 131072,
    "qwen-max-latest": 64000,
}
# 更新所有可⽤模型的字典，添加DeepSeeK模型
ALL_AVAILABLE_MODELS.update(LLM_MODELS_CONF)
# 更新聊天模型的字典，添加DeepSeeK模型
CHAT_MODELS.update(LLM_MODELS_CONF)

# 设置密集向量嵌入模型,模型维度要跟milvus建表时候的向量维度保持一致
embedding = HuggingFaceEmbedding(
    # model_name="BAAI/bge-small-en-v1.5" #384维度
    model_name="BAAI/bge-base-en-v1.5"  # 768维度
)
dim = 768
Settings.embed_model = embedding

from llama_index.llms.openai import OpenAI

llm_url = "http://localhost:11434/v1"
llm_model = "qwen2.5-coder:7b"

llm_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
llm_model = "qwen-max-latest"
# 阿里云 API-KEY, 需要设置环境变量ALIYUN-API-KEY
# 因为调用本机的ollama模型,资源消耗太大,容易卡,还会导致 milvus 数据库宕机
llm_api_key = os.environ['ALIYUN-API-KEY'] or 'aa'

llm = OpenAI(model=llm_model, api_base=llm_url, api_key=llm_api_key, temperature=0.7)
Settings.llm =llm

#官方用的默认稀疏向量检索函数
from llama_index.vector_stores.milvus.utils import get_default_sparse_embedding_function

sparse_embedding_function = get_default_sparse_embedding_function()

collection_name = 'knowledge_base'
# collection_name = 'limd_li'
milvus_url = "http://localhost:19530"
milvus_token = 'root:Milvus'

milvus_url=os.environ['ZILLIZ-MILVUS-URL']
milvus_token=os.environ['ZILLIZ-MILVUS-TOKEN']

input_dir = project_setting.knowledge_files_path
input_files = [
    input_dir + '/knowledge.txt',
    input_dir + '/paul_graham_essay.txt',
]


class LlamaIndexDemo5b:
    vector_store: MilvusVectorStore = None
    index = None
    fold_index = None


base = LlamaIndexDemo5b()


def init_milvus():
    base.vector_store = MilvusVectorStore(
        dim=dim,
        uri=milvus_url,
        token=milvus_token,
        collection_name=collection_name,
        overwrite=False,
        hybrid_ranker="RRFRanker",
        hybrid_ranker_params={"k": 60},
        # enable_sparse=True,
        # sparse_embedding_function=sparse_embedding_function
    )




def drop_collection():
    if base.vector_store.client.has_collection(collection_name) is True:
        """测试的,存在则先删除集合再创建"""
        base.vector_store.client.drop_collection(collection_name)


def init_index():
    # storage_context = StorageContext.from_defaults(vector_store=base.vector_store, persist_dir=DEFAULT_PERSIST_DIR, )
    storage_context = StorageContext.from_defaults(vector_store=base.vector_store)
    docs = SimpleDirectoryReader(input_files=input_files).load_data()
    args = {
        'summary': '索引'
    }
    base.index = VectorStoreIndex.from_documents(
        documents=docs, storage_context=storage_context, show_progress=True, **args)
    return base.index


def init_index2():
    # 获取索引
    base.index = VectorStoreIndex.from_vector_store(base.vector_store)
    return base.index

doc_ids=[]


def format_doc_id( kb_doc_id):
    return f'kb_doc_{kb_doc_id}'


def add_milvus_doc():
    kb_doc_id=20
    for input_file in input_files:
        kb_doc_id+=1
        docs = SimpleDirectoryReader(input_files=[input_file],filename_as_id=True).load_data()
        _index: VectorStoreIndex = base.index
        for doc in docs:
            # store_doc_id = format_doc_id(kb_doc_id)
            # doc.doc_id=store_doc_id
            #自定义将 文件名作为doc_id,有利于后续删除操作
            doc.doc_id=format_filename(input_file)
            print(str_prominent(doc.doc_id))
            doc_ids.append(doc.doc_id)
            _index.insert(document=doc)
        print(str_prominent(doc_ids))
    # Document()
    # _index.insert(document=Document)

def del_milvus_doc():
    _index:VectorStoreIndex=base.index
    for input_file in input_files:
        doc_id=format_filename(input_file)
        print(str_prominent(f'delete {doc_id}'))
        _index.delete(doc_id=doc_id)
    # _index.delete_ref_doc(ref_doc_id=format_doc_id(21))

def reactagent():
    # 构建检索器
    from llama_index.core.retrievers import VectorIndexRetriever
    # 想要自定义参数，可以构造参数字典
    kwargs = {'similarity_top_k': 5, 'index': base.index, 'dimensions': dim}  # 必要参数
    retriever = VectorIndexRetriever(**kwargs)


    # 构建合成器
    from llama_index.core.response_synthesizers import get_response_synthesizer
    response_synthesizer = get_response_synthesizer(llm=llm, streaming=True)

    # 构建问答引擎
    from llama_index.core.query_engine import RetrieverQueryEngine
    engine = RetrieverQueryEngine(
        retriever=retriever,
        response_synthesizer=response_synthesizer,
    )
    # 配置查询工具
    from llama_index.core.tools import QueryEngineTool
    from llama_index.core.tools import ToolMetadata
    query_engine_tools = [
        QueryEngineTool(
            query_engine=engine,
            metadata=ToolMetadata(
                name="RAG工具",
                description=(
                    "用于在原文中检索相关信息"
                ),
            ),
        ),
    ]
    # 创建ReAct Agent
    from llama_index.core.agent import ReActAgent
    agent = ReActAgent.from_tools(query_engine_tools, llm=llm, verbose=True)
    query="What are the applications of Agent AI systems ?"
    query="广州大学"
    res = agent.chat(query)
    msg_sout(f'{"-" * 30} 问题 {"-" * 30} ')
    msg_sout(query)
    msg_sout(f'{"-" * 30} 答案 {"-" * 30} ')
    msg_sout(f"{str(res)}\n")


#test start
from llama_index.core.vector_stores import VectorStoreQuery
from llama_index.core.vector_stores.types import VectorStoreQueryMode


def to_query(query):
    query_embedding = embedding.get_query_embedding(query)
    # res = base.vector_store.query(VectorStoreQuery(query_str=query, query_embedding=query_embedding, mode=VectorStoreQueryMode.HYBRID))
    res = base.vector_store.query(
        VectorStoreQuery(query_str=query, query_embedding=query_embedding, mode=VectorStoreQueryMode.DEFAULT))
    # res = base.vector_store.query(VectorStoreQuery(query_str=query))
    return res


#test end

def to_query2(index, query):
    # 创建查询引擎
    query_engine = index.as_query_engine(vector_store_query_mode=VectorStoreQueryMode.DEFAULT)
    # 执行查询
    res = query_engine.query(query)
    return res


# 创建主窗口
root = tk.Tk()
root.title("index持久化再加载测试")
root.geometry("1300x900")


def msg_sout(msg):
    # text_box.insert(tk.END, msg)
    text_box.insert(tk.END, f"{msg}\n")
    text_box.see(tk.END)
    # print('msg_sout')


# 定义按钮列表以便于禁用和启用
buttons = []


def disable_buttons():
    for button in buttons:
        button.config(state=tk.DISABLED)


def enable_buttons():
    for button in buttons:
        button.config(state=tk.NORMAL)


index_save_path = 'index_cache'


def button_init_milvus():
    init_milvus()
def button_delete_fold_index():
    """"""
    from llama_index.core.storage.docstore import SimpleDocumentStore
    tmp:SimpleDocumentStore=base.fold_index.docstore
    ids=[]
    for key in tmp.docs:
        from llama_index.core.schema import BaseNode
        doc:BaseNode =tmp.docs[key]
        ids.append(doc.ref_doc_id)
        print(doc)

    for ref_doc_id in ids:
        base.fold_index.delete_ref_doc(ref_doc_id,delete_from_docstore=True)
        print(f'doc_id={ref_doc_id} 删除成功')

    # 索引持久化到磁盘 默认目录 persist_dir='./storage'
    base.fold_index.storage_context.persist(persist_dir=DEFAULT_PERSIST_DIR)


def button_create_fold_index():
    docs = SimpleDirectoryReader(input_files=input_files).load_data()
    # 创建新索引，显示 embedding 进度
    index = VectorStoreIndex.from_documents(docs, show_progress=True)
    # 索引持久化到磁盘 默认目录 persist_dir='./storage'
    index.storage_context.persist(persist_dir=DEFAULT_PERSIST_DIR)
    base.fold_index = index

    msg_sout("索引持久化成功\n")

def button_update_fold_index():
    docs = SimpleDirectoryReader(input_files=input_files).load_data()
    index:VectorStoreIndex=base.fold_index
    for doc in docs:
        index.update_ref_doc(doc)
    index.storage_context.persist(persist_dir=DEFAULT_PERSIST_DIR)
    msg_sout("索引更新持久化成功\n")

def button_load_fold_index():
    # 在下次启动时，使用相同的 StorageContext 来加载索引

    msg_sout("正在加载现有索引。..")
    storage_context = StorageContext.from_defaults(persist_dir=DEFAULT_PERSIST_DIR)
    from llama_index.core import load_index_from_storage
    structs = storage_context.index_store.index_structs()
    # print(structs)
    index_id=None
    for key in structs:
        # print(key)
        # print(key.index_id)
        index_id = key.index_id
    base.fold_index = load_index_from_storage(storage_context,index_id=index_id)
    msg_sout("读取持久化后的索引成功\n")

def button_query_fold_index():
    query = "广州大学"
    res = to_query2(base.fold_index, query)
    msg_sout(f'{"-" * 30} 问题 {"-" * 30} \n')
    msg_sout(query + "\n")
    msg_sout(f'{"-" * 30} 答案 {"-" * 30} \n')
    msg_sout(str(res) + "\n")


def button_drop_collection():
    drop_collection()
    msg_sout("删除集合成功\n")


def button_create_collection():
    init_milvus()
    msg_sout("集合创建成功\n")


def button_init_index1():
    base.index = init_index()
    msg_sout("初始化索引成功\n")


def button_init_index2():
    base.index = init_index2()
    msg_sout("初始化索引成功\n")

def button_add_milvus_doc():
    add_milvus_doc()
    msg_sout("添加文档向量成功\n")

def button_del_milvus_doc():
    del_milvus_doc()
    msg_sout("删除文档向量成功\n")

def button_reactagent():
    msg_sout("知识库作为工具供AI使用\n")
    reactagent()


def button_query01():
    query = "广州大学"
    res = to_query(query)
    msg_sout(f'{"-" * 30} 问题 {"-" * 30} \n')
    msg_sout(query + "\n")
    msg_sout(f'{"-" * 30} 答案 {"-" * 30} \n')
    msg_sout(str(res) + "\n")


def button_query02():
    query = "广州大学"
    res = to_query2(base.index, query)
    msg_sout(f'{"-" * 30} 问题 {"-" * 30} \n')
    msg_sout(query + "\n")
    msg_sout(f'{"-" * 30} 答案 {"-" * 30} \n')
    msg_sout(str(res) + "\n")


#处理线程,避免UI卡死
action_thread = None


def button_async_fun(_fun):
    disable_buttons()
    root.update()  # 更新主界面以显示进度

    def async_fun(_fun):
        try:
            _fun()
            """"""
        except Exception as e:
            text_box.insert(tk.END, f"执行失败: {str(e)}\n")
            # 打印完整的异常信息
            import traceback
            traceback.print_exc(file=sys.stderr)
        enable_buttons()

    global action_thread
    if action_thread is None or not action_thread.is_alive():
        import threading
        action_thread = threading.Thread(target=async_fun, args=(_fun,))
        action_thread.start()
        # 等待线程结束,调用这个函数会阻塞线程,直到线程结束
        # action_thread.join()
        # enable_buttons()
    else:
        enable_buttons()
        text_box.showwarning("提示", "请等待上一个操作完成")


def button_async_fun2(_fun):
    disable_buttons()
    root.update()  # 更新主界面以显示进度

    async def async_fun2(_fun):
        try:
            _fun()
            """"""
        except Exception as e:
            text_box.insert(tk.END, f"执行失败: {str(e)}\n")
            # 打印完整的异常信息
            import traceback
            traceback.print_exc(file=sys.stderr)
        enable_buttons()

    import asyncio
    asyncio.run(async_fun2(_fun))


# 创建按钮框架
button_frame = ttk.Frame(root)
button_frame.pack(pady=10)

label = ttk.Label(button_frame, text="milvus 存储")
label.pack(side=tk.LEFT)

init_index_button = ttk.Button(button_frame, text="初始化milvus", command=lambda: button_async_fun2(button_init_milvus))
init_index_button.pack(side=tk.LEFT)
buttons.append(init_index_button)

init_index_button = ttk.Button(button_frame, text="删除集合", command=lambda: button_async_fun(button_drop_collection))
init_index_button.pack(side=tk.LEFT)
buttons.append(init_index_button)

init_index_button = ttk.Button(button_frame, text="创建集合", command=lambda: button_async_fun2(button_create_collection))
init_index_button.pack(side=tk.LEFT)
buttons.append(init_index_button)

# 添加初始化索引按钮
init_index_button = ttk.Button(button_frame, text="创建索引", command=lambda: button_async_fun(button_init_index1))
init_index_button.pack(side=tk.LEFT)
buttons.append(init_index_button)

init_index_button = ttk.Button(button_frame, text="加载milvus索引",
                               command=lambda: button_async_fun(button_init_index2))
init_index_button.pack(side=tk.LEFT)
buttons.append(init_index_button)

# 添加查询按钮
query_button = ttk.Button(button_frame, text="查询1", command=lambda: button_async_fun(button_query01))
query_button.pack(side=tk.LEFT)
buttons.append(query_button)

query_button = ttk.Button(button_frame, text="查询2", command=lambda: button_async_fun(button_query02))
query_button.pack(side=tk.LEFT)
buttons.append(query_button)

# 创建按钮框架
button_frame = ttk.Frame(root)
button_frame.pack(pady=10)

label = ttk.Label(button_frame, text="milvus 存储")
label.pack(side=tk.LEFT)

query_button = ttk.Button(button_frame, text="新增文档向量", command=lambda: button_async_fun(button_add_milvus_doc))
query_button.pack(side=tk.LEFT)
buttons.append(query_button)

query_button = ttk.Button(button_frame, text="删除文档向量", command=lambda: button_async_fun(button_del_milvus_doc))
query_button.pack(side=tk.LEFT)
buttons.append(query_button)

# 创建按钮框架
button_frame = ttk.Frame(root)
button_frame.pack(pady=10)

label = ttk.Label(button_frame, text="milvus 存储 作为 ReActAgent的工具")
label.pack(side=tk.LEFT)

query_button = ttk.Button(button_frame, text="查询", command=lambda: button_async_fun(button_reactagent))
query_button.pack(side=tk.LEFT)
buttons.append(query_button)

# 创建按钮框架
button_frame = ttk.Frame(root)
button_frame.pack(pady=10)

label = ttk.Label(button_frame, text="目录 存储")
label.pack(side=tk.LEFT)



query_button = ttk.Button(button_frame, text="删除文档索引", command=lambda: button_async_fun(button_delete_fold_index))
query_button.pack(side=tk.LEFT)
buttons.append(query_button)

query_button = ttk.Button(button_frame, text="创建索引", command=lambda: button_async_fun(button_create_fold_index))
query_button.pack(side=tk.LEFT)
buttons.append(query_button)

query_button = ttk.Button(button_frame, text="更新索引(删除文档索引后操作)", command=lambda: button_async_fun(button_update_fold_index))
query_button.pack(side=tk.LEFT)
buttons.append(query_button)


query_button = ttk.Button(button_frame, text="加载索引", command=lambda: button_async_fun(button_load_fold_index))
query_button.pack(side=tk.LEFT)
buttons.append(query_button)

query_button = ttk.Button(button_frame, text="查询", command=lambda: button_async_fun(button_query_fold_index))
query_button.pack(side=tk.LEFT)
buttons.append(query_button)

# 创建外框架以容纳纵向滚动条和另一个外框架
frame = ttk.Frame(root)
frame.pack(fill=tk.BOTH, expand=True)

# 创建外框架以容纳横向滚动条和文本框
panel = ttk.Frame(frame)
panel.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

# 创建多行文本框
text_box = tk.Text(panel, wrap=tk.NONE)
text_box.pack(side=tk.TOP, fill=tk.BOTH, expand=True)

# 创建滚动条并将其与文本框关联
scrollbar_y = tk.Scrollbar(frame, command=text_box.yview)
scrollbar_x = ttk.Scrollbar(panel, orient=tk.HORIZONTAL, command=text_box.xview)

scrollbar_y.pack(side=tk.RIGHT, fill=tk.Y)
scrollbar_x.pack(side=tk.BOTTOM, fill=tk.X)

text_box.config(yscrollcommand=scrollbar_y.set, xscrollcommand=scrollbar_x.set)

readme = """
# milvus方式 测试
## 第一次启动
初始化milvus -> 删除集合 -> 创建集合 -> 创建索引 -> 查询1 -> 查询2   

## 第二次启动
初始化milvus -> 加载 milvus索引  ->  查询1 -> 查询2


# milvus方式 测试2,增加文档,删除文档
初始化milvus -> 删除集合 -> 创建集合 -> 加载 milvus索引 -> 新增文档向量 -> 查询1 -> 查询2   -> 删除文档向量 -> 查询1 -> 查询2  新增文档向量  -> 查询1 -> 查询2


# milvus 存储 作为 ReActAgent的工具
初始化milvus -> 删除集合 -> 创建集合 -> 加载 milvus索引 -> 新增文档向量 -> 查询

# 目录方式 测试 
索引存储目录是 ./storage ,如果测试过程中顺序错了,结果不对,可先删除该目录后再试
## 第一次启动
创建索引 -> 查询
    
## 第二次启动
加载索引 -> 查询

## 第三次启动
加载索引 -> 查询 -> 删除索引 -> 查询(无数据)  -> 更新索引 -> 查询

    
"""

msg_sout(readme)

# 运行主循环
root.mainloop()
